
#include "STP.h"
#include "bridge.h"
#include "log.h"

void RunStateMachines (BRIDGE* bridge);
static unsigned long GetInstanceCountForStateMachine (const SM_INFO* smInfo, int portCount, int treeCount);
static void RecomputePrioritiesAndPortRoles (BRIDGE* bridge, int treeIndex);

static const unsigned char NullBridgeAddress [6] = { 0, 0, 0, 0, 0, 0 };

// ============================================================================
// numberOfMstis must be 2..64, as per 5.3.1.1.2) in 802.1Q-2005

BRIDGE* STP_CreateBridge (int portCount, int treeCount, const STP_CALLBACKS* callbacks, STP_MODE compatibilityMode, unsigned long debugLogBufferSize)
{
	BRIDGE* bridge = (BRIDGE*) callbacks->allocAndZeroMemory (sizeof (BRIDGE));

	assert (debugLogBufferSize >= 2); // one byte for the data, one for the null terminator of the string passed to the callback
	
	bridge->compatibilityMode = compatibilityMode;
	switch (compatibilityMode)
	{
		case STP_MODE_802_1w_2001:
			assert (treeCount == 1);
			bridge->ForceProtocolVersion = 2;
			bridge->smInterface = &smInterface_802_1w_2001;
			break;

		case STP_MODE_802_1D_2004:
			assert (treeCount == 1);
			bridge->ForceProtocolVersion = 2;
			bridge->smInterface = &smInterface_802_1D_2004;
			break;

		case STP_MODE_802_1Q_2005:
			assert ((treeCount >= 1) && (treeCount <= 64));
			bridge->ForceProtocolVersion = 3;
			bridge->smInterface = &smInterface_802_1Q_2005;
			break;

		default:
			assert (false);
	}

	bridge->callbacks = callbacks;
	bridge->portCount = portCount;
	bridge->treeCount = treeCount;

	bridge->logBuffer = (char*) callbacks->allocAndZeroMemory (debugLogBufferSize);
	bridge->logBufferSize = debugLogBufferSize;
	bridge->logBufferUsedSize = 0;
	bridge->logCurrentPort = -1;
	bridge->logCurrentTree = -1;

	// ------------------------------------------------------------------------
	// alloc space for state array
	unsigned long stateMachineInstanceCount = 0;
	for (unsigned long i = 0; i < bridge->smInterface->smInfoCount; i++)
		stateMachineInstanceCount += GetInstanceCountForStateMachine (&bridge->smInterface->smInfo [i], portCount, treeCount);

	stateMachineInstanceCount += GetInstanceCountForStateMachine (bridge->smInterface->transmitSmInfo, portCount, treeCount);

	bridge->states = (SM_STATE*) callbacks->allocAndZeroMemory (stateMachineInstanceCount * sizeof (SM_STATE));

	// ------------------------------------------------------------------------

	bridge->trees = (BRIDGE_TREE**) callbacks->allocAndZeroMemory (treeCount * sizeof (BRIDGE_TREE*));
	bridge->ports = (PORT**) callbacks->allocAndZeroMemory (portCount * sizeof (PORT*));
	
	// per-bridge CIST vars
	bridge->trees [CIST_INDEX] = (BRIDGE_TREE*) callbacks->allocAndZeroMemory (sizeof (BRIDGE_TREE));
	bridge->trees [CIST_INDEX]->SetBridgeIdentifier (0x8000, CIST_INDEX, NullBridgeAddress);
	bridge->trees [CIST_INDEX]->BridgeTimes.HelloTime		= bridge->HelloTime;
	bridge->trees [CIST_INDEX]->BridgeTimes.remainingHops	= bridge->MaxHops;
	bridge->trees [CIST_INDEX]->BridgeTimes.ForwardDelay	= bridge->BridgeForwardDelay;
	bridge->trees [CIST_INDEX]->BridgeTimes.MaxAge			= bridge->BridgeMaxAge;
	bridge->trees [CIST_INDEX]->BridgeTimes.MessageAge		= 0;

	// per-bridge MSTI vars
	for (int treeIndex = 1; treeIndex < treeCount; treeIndex++)
	{
		bridge->trees [treeIndex] = (BRIDGE_TREE*) callbacks->allocAndZeroMemory (sizeof (BRIDGE_TREE));
		bridge->trees [treeIndex]->SetBridgeIdentifier (0x8000, treeIndex, NullBridgeAddress);
	}

	// per-port vars
	for (int portIndex = 0; portIndex < bridge->portCount; portIndex++)
	{
		bridge->ports [portIndex] = (PORT*) callbacks->allocAndZeroMemory (sizeof (PORT));
		PORT* port = bridge->ports [portIndex];
		
		port->trees = (PORT_TREE**) callbacks->allocAndZeroMemory (treeCount * sizeof (PORT_TREE*));
			
		// per-port CIST and MSTI vars
		for (int treeIndex = 0; treeIndex < treeCount; treeIndex++)
		{
			port->trees [treeIndex] = (PORT_TREE*) callbacks->allocAndZeroMemory (sizeof (PORT_TREE));
			port->trees [treeIndex]->portId = 0x8000 | (portIndex + 1);
			port->trees [treeIndex]->portTimes.HelloTime = bridge->HelloTime;
		}

		// give the port a default name
		sprintf (port->debugName, "%04x", 0x8000 | (portIndex + 1));
		
		port->AutoEdge = 1;
	}
	
	bridge->receivedBpduContent = (const MSTP_BPDU*) 0xcccccccc; // to cause an exception on access
		
	bridge->MstConfigId.ConfigurationIdentifierFormatSelector = 0;
	memset (bridge->MstConfigId.ConfigurationName, 0, 32);
	bridge->MstConfigId.RevisionLevel [0] = 0;
	bridge->MstConfigId.RevisionLevel [0] = 0;
	memcpy (bridge->MstConfigId.ConfigurationDigest, "\xAC\x36\x17\x7F\x50\x28\x3C\xD4\xB8\x38\x21\xD8\xAB\x26\xDE\x62", 16); // All VIDs map to the CIST, no VID mapped to any MSTI

	return bridge;
}

// ============================================================================

void STP_DestroyBridge (BRIDGE* bridge)
{
	assert (bridge->started == false);

	for (int portIndex = 0; portIndex < bridge->portCount; portIndex++)
	{
		for (int treeIndex = 0; treeIndex < bridge->treeCount; treeIndex++)
			bridge->callbacks->freeMemory (bridge->ports [portIndex]->trees [treeIndex]);

		bridge->callbacks->freeMemory (bridge->ports [portIndex]->trees);
		
		bridge->callbacks->freeMemory (bridge->ports [portIndex]);
	}

	for (int treeIndex = 0; treeIndex < bridge->treeCount; treeIndex++)
		bridge->callbacks->freeMemory (bridge->trees [treeIndex]);

	bridge->callbacks->freeMemory (bridge->ports);
	bridge->callbacks->freeMemory (bridge->trees);
	bridge->callbacks->freeMemory (bridge->states);
	bridge->callbacks->freeMemory (bridge->logBuffer);

	bridge->callbacks->freeMemory (bridge);
}

// ============================================================================

void STP_StartBridge (BRIDGE* bridge, unsigned long timestampMilliseconds)
{
	LOG (bridge, -1, -1, "{T}: Starting the bridge...\r\n", timestampMilliseconds);

	assert (bridge->started == false);
	
	// check that the MAC address was set
	for (int treeIndex = 0; treeIndex < bridge->treeCount; treeIndex++)
		assert (bridge->trees [treeIndex]->GetBridgeIdentifier ()._address.IsZero () == false);
	
	bridge->started = true;
	
	bridge->BEGIN = true;
	RunStateMachines (bridge);
	bridge->BEGIN = false;	

	LOG (bridge, -1, -1, "Bridge started.\r\n");
	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

// ============================================================================

void STP_StopBridge (BRIDGE* bridge, unsigned long timestampMilliseconds)
{
	LOG (bridge, -1, -1, "{T}: Bridge stopped.\r\n", timestampMilliseconds);

	bridge->started = false;

	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

// ============================================================================

void STP_OnBridgeAddressChanged (BRIDGE* bridge, const unsigned char address [6], unsigned long timestampMilliseconds)
{
	BRIDGE_ADDRESS ba;
	ba.SetValue (address);
	LOG (bridge, -1, -1, "{T}: MAC address changed to {BA}\r\n", timestampMilliseconds, &ba);
	
	for (int treeIndex = 0; treeIndex < bridge->treeCount; treeIndex++)
	{
		// change the MAC address without changing the priority
		BRIDGE_ID bid = bridge->trees [treeIndex]->GetBridgeIdentifier ();
		bid.SetAddress (address);
		bridge->trees [treeIndex]->SetBridgeIdentifier (bid);
	}
	
	if (bridge->started)
		RecomputePrioritiesAndPortRoles (bridge, CIST_INDEX); // we pass CIST_INDEX as this will cause the function to recompute for all trees

	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

// ============================================================================

void STP_OnPortEnabled (BRIDGE* bridge, int portIndex, unsigned long speedMegabitsPerSecond, DUPLEX duplex, unsigned long timestampMilliseconds)
{
	LOG (bridge, -1, -1, "{T}: Port {S} good\r\n", timestampMilliseconds, bridge->ports [portIndex]->debugName);

	assert (bridge->ports [portIndex]->portEnabled == false);

	bridge->ports [portIndex]->portEnabled = true;

	if (bridge->ports [portIndex]->adminPointToPointMAC)
		bridge->ports [portIndex]->operPointToPointMAC = true;
	else
		bridge->ports [portIndex]->operPointToPointMAC = (duplex == FULL_DUPLEX);
	
	if (bridge->started)
		RunStateMachines (bridge);

	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

// ============================================================================

void STP_OnPortDisabled (BRIDGE* bridge, int portIndex, unsigned long timestampMilliseconds)
{
	LOG (bridge, -1, -1, "{T}: Port {S} down\r\n", timestampMilliseconds, bridge->ports [portIndex]->debugName);

	assert (bridge->ports [portIndex]->portEnabled);

	bridge->ports [portIndex]->portEnabled = false;

	if (bridge->started)
		RunStateMachines (bridge);

	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

// ============================================================================

void STP_OnOneSecondTick (BRIDGE* bridge, unsigned long timestampMilliseconds)
{
	if (bridge->started)
	{
		LOG (bridge, -1, -1, "{T}: One second:\r\n", timestampMilliseconds);
	
		for (int givenPort = 0; givenPort < bridge->portCount; givenPort++)
			bridge->ports [givenPort]->tick = true;

		RunStateMachines (bridge);

		LOG (bridge, -1, -1, "------------------------------------\r\n");
		FLUSH_LOG (bridge);
	}
}

// ============================================================================

static void ProcessStpConfigBpdu (BRIDGE* bridge, int portIndex, const unsigned char* bpdu, int bpduSize)
{
	assert (false);
}

static void ProcessStpTcnBpdu (BRIDGE* bridge, int portIndex, const unsigned char* bpdu, int bpduSize)
{
	assert (false);
}

static void ProcessMstOrRstBpdu (BRIDGE* bridge, int portIndex, VALIDATED_BPDU_TYPE type, const MSTP_BPDU* bpdu, int bpduSize)
{
	assert ((unsigned long) bridge->receivedBpduContent == 0xcccccccc);
	assert (bridge->receivedBpduType == VALIDATED_BPDU_TYPE_UNKNOWN);
	assert (bridge->ports [portIndex]->rcvdBpdu == false);
	
	bridge->receivedBpduContent = bpdu;
	bridge->receivedBpduType = type;
	bridge->ports [portIndex]->rcvdBpdu = true;
	
	RunStateMachines (bridge);
	
	bridge->receivedBpduContent = (const MSTP_BPDU*) 0xcccccccc; // to cause an exception on access
	bridge->receivedBpduType = VALIDATED_BPDU_TYPE_UNKNOWN; // to cause asserts on access

	// If the port is enabled, check that the state machines did process the BPDU
	if (bridge->ports [portIndex]->portEnabled)
	{
		assert (bridge->ports [portIndex]->rcvdBpdu == false);
	}
	else
		bridge->ports [portIndex]->rcvdBpdu = false;
}

// ============================================================================

void STP_OnBpduReceived (BRIDGE* bridge, int portIndex, const unsigned char* bpdu, int bpduSize, unsigned long timestampMilliseconds)
{
	LOG (bridge, -1, -1, "{T}: BPDU received on Port {S}:\r\n", timestampMilliseconds, bridge->ports [portIndex]->debugName);
	
	if (bridge->started == false)
	{
		LOG (bridge, -1, -1, "Bridge currently stopped.\r\n");
	}
	else
	{
		enum VALIDATED_BPDU_TYPE type = MSTP_GetValidatedBpduType (bpdu, bpduSize);
		switch (type)
		{
			case VALIDATED_BPDU_TYPE_STP_CONFIG:
				ProcessStpConfigBpdu (bridge, portIndex, bpdu, bpduSize);
				break;
				
			case VALIDATED_BPDU_TYPE_STP_TCN:
				ProcessStpTcnBpdu (bridge, portIndex, bpdu, bpduSize);
				break;
				
			case VALIDATED_BPDU_TYPE_RST:
				LOG (bridge, portIndex, -1, "RSTP BPDU:\r\n");
				DumpRstpBpdu (bridge, portIndex, -1, (const MSTP_BPDU*) bpdu, 2);
				ProcessMstOrRstBpdu (bridge, portIndex, VALIDATED_BPDU_TYPE_RST, (const MSTP_BPDU*) bpdu, bpduSize);
				break;
				
			case VALIDATED_BPDU_TYPE_MST:
				LOG (bridge, portIndex, -1, "MSTP BPDU:\r\n");
				DumpMstpBpdu (bridge, portIndex, -1, (const MSTP_BPDU*) bpdu, 2);
				ProcessMstOrRstBpdu (bridge, portIndex, VALIDATED_BPDU_TYPE_MST, (const MSTP_BPDU*) bpdu, bpduSize);
				break;
				
			default:
				LOG (bridge, portIndex, -1, "Invalid BPDU received!\r\n");
		}	
	}

	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

// ============================================================================

bool STP_IsBridgeStarted (BRIDGE* bridge)
{
	return bridge->started;
}

// ============================================================================

void STP_SetPortName (BRIDGE* bridge, int portIndex, const char* name)
{
	size_t nameLen = strlen (name);
	assert (nameLen <= PORT_DEBUG_NAME_MAX_LENGTH);

	strncpy (bridge->ports [portIndex]->debugName, name, PORT_DEBUG_NAME_MAX_LENGTH);
	bridge->ports [portIndex]->debugName [nameLen] = 0;
}

// ============================================================================

const char* STP_GetPortName (BRIDGE* bridge, int portIndex)
{
	return bridge->ports [portIndex]->debugName;
}

// ============================================================================

void STP_EnableLogging (BRIDGE* bridge, bool enable)
{
	bridge->loggingEnabled = enable;
}

// ============================================================================

bool STP_IsLoggingEnabled (BRIDGE* bridge)
{
	return bridge->loggingEnabled;
}

// ============================================================================

bool RunStateMachineInstance (BRIDGE* bridge, const SM_INFO* smInfo, int givenPort, int givenTree, SM_STATE* statePtr)
{
	bool changed = false;

rep:
	SM_STATE newState = smInfo->checkConditions (bridge, givenPort, givenTree, *statePtr);
	if (newState != 0)
	{
		if (givenPort == -1)
			LOG (bridge, givenPort, givenTree, "Bridge: ");
		else
			LOG (bridge, givenPort, givenTree, "{S}: ", bridge->ports [givenPort]->debugName);
		
		if (bridge->compatibilityMode == STP_MODE_802_1Q_2005)
		{
			if (givenTree == CIST_INDEX)
				LOG (bridge, givenPort, givenTree, "CIST: ");
			else if (givenTree > 0)
				LOG (bridge, givenPort, givenTree, "MSTI_{D}: ", givenTree);
		}
		
		const char* currentStateName = smInfo->getStateName (*statePtr);
		const char* newStateName = smInfo->getStateName (newState);

		//LOG (bridge, givenPort, givenTree, "{S}: {S} -> {S}\r\n", smInfo->smName, currentStateName, newStateName);
		LOG (bridge, givenPort, givenTree, "{S}: -> {S}\r\n", smInfo->smName, newStateName);

		smInfo->initState (bridge, givenPort, givenTree, newState);

		*statePtr = newState;
		changed = true;
		goto rep;
	}

	return changed;
}

// ============================================================================

bool RunStateMachineInstances (BRIDGE* bridge, const SM_INFO* smInfo, SM_STATE** statePtr)
{
	volatile bool changed = false;

	switch (smInfo->instanceType)
	{
		case SM_INFO::PER_BRIDGE:
			changed |= RunStateMachineInstance (bridge, smInfo, -1, -1, *statePtr);
			(*statePtr)++;
			break;

		case SM_INFO::PER_BRIDGE_PER_TREE:
			for (int treeIndex = 0; treeIndex < bridge->treeCount; treeIndex++)
			{
				changed |= RunStateMachineInstance (bridge, smInfo, -1, treeIndex, *statePtr);
				(*statePtr)++;
			}
			break;

		case SM_INFO::PER_PORT:
			for (int portIndex = 0; portIndex < bridge->portCount; portIndex++)
			{
				changed |= RunStateMachineInstance (bridge, smInfo, portIndex, -1, *statePtr);
				(*statePtr)++;
			}
			break;

		case SM_INFO::PER_PORT_PER_TREE:
			for (int treeIndex = 0; treeIndex < bridge->treeCount; treeIndex++)
			{
				for (int portIndex = 0; portIndex < bridge->portCount; portIndex++)
				{
					changed |= RunStateMachineInstance (bridge, smInfo, portIndex, treeIndex, *statePtr);
					(*statePtr)++;
				}
			}
			break;
	}

	return changed;
}

// ============================================================================

void RunStateMachines (BRIDGE* bridge)
{
	volatile bool changed;
	
	do
	{
		SM_STATE* statePtr = bridge->states;

		changed = false;
		
		for (unsigned long i = 0; i < bridge->smInterface->smInfoCount; i++)
			changed |= RunStateMachineInstances (bridge, &bridge->smInterface->smInfo [i], &statePtr);
		
		if (changed == false)
			changed |= RunStateMachineInstances (bridge, bridge->smInterface->transmitSmInfo, &statePtr);
		
	} while (changed);
}

// ============================================================================

static unsigned long GetInstanceCountForStateMachine (const SM_INFO* smInfo, int portCount, int treeCount)
{
	switch (smInfo->instanceType)
	{
		case SM_INFO::PER_BRIDGE:			return 1;
		case SM_INFO::PER_PORT:				return portCount;
		case SM_INFO::PER_BRIDGE_PER_TREE:	return treeCount;
		case SM_INFO::PER_PORT_PER_TREE:	return treeCount * portCount;
		default:
			assert (false);
			return 0;
	}
}

// ============================================================================

void STP_SetPortAdminEdge (BRIDGE* bridge, int portIndex, bool newAdminEdge)
{
	bridge->ports [portIndex]->AdminEdge = newAdminEdge;
}

// ============================================================================

void STP_SetPortAutoEdge (BRIDGE* bridge, int portIndex, bool newAutoEdge)
{
	bridge->ports [portIndex]->AutoEdge = newAutoEdge;
}

// ============================================================================

void STP_SetPortAdminP2P (BRIDGE* bridge, int portIndex, bool newAdminP2P)
{
	bridge->ports [portIndex]->adminPointToPointMAC = newAdminP2P;
}

// ============================================================================

static void RecomputePrioritiesAndPortRoles (BRIDGE* bridge, int treeIndex)
{
	// From page 357 of 802.1Q-2011:
	// BridgeIdentifier, BridgePriority, and BridgeTimes are not modified by the operation of the spanning tree
	// protocols but are treated as constants by the state machines. If they are modified by management, spanning
	// tree priority vectors and Port Role assignments for shall be recomputed, as specified by the operation of the
	// Port Role Selection state machine (13.34) by clearing selected (13.25) and setting reselect (13.25) for all
	// Bridge Ports for the relevant MSTI and for all trees if the CIST parameter is changed.

	if (treeIndex == CIST_INDEX)
	{
		// Recompute all trees.
		// Note that callers of this function expect recomputation for all trees when CIST_INDEX is passed, so don't change this functionality.
		for (int ti = 0; ti < bridge->treeCount; ti++)
		{
			for (int pi = 0; pi < bridge->portCount; pi++)
			{
				bridge->ports [pi]->trees [ti]->selected = false;
				bridge->ports [pi]->trees [ti]->reselect = true;
			}
		}
	}
	else
	{
		// recompute specified MSTI
		assert (false); // not implemented
	}
	
	RunStateMachines (bridge);
}

void STP_SetBridgePriority (BRIDGE* bridge, int treeIndex, unsigned short bridgePriority, unsigned long timestampMilliseconds)
{
	// See table 13-3 on page 348 of 802.1Q-2011
		
	assert (bridgePriority <= 61440);
	assert ((bridgePriority % 4096) == 0);
	
	assert ((treeIndex >= 0) && (treeIndex < bridge->treeCount));

	LOG (bridge, -1, -1, "{T}: Setting bridge priority: tree {TN} prio = {D}...\r\n",
		 timestampMilliseconds,
		 treeIndex,
		 bridgePriority);

	BRIDGE_ID bid = bridge->trees [treeIndex]->GetBridgeIdentifier ();
	bid.SetPriority (bridgePriority, treeIndex);
	bridge->trees [treeIndex]->SetBridgeIdentifier (bid);

	if (bridge->started)
		RecomputePrioritiesAndPortRoles (bridge, treeIndex);

	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

unsigned short STP_GetBridgePriority (BRIDGE* bridge, int treeIndex)
{
	assert ((treeIndex >= 0) && (treeIndex < bridge->treeCount));
	
	return bridge->trees [treeIndex]->GetBridgeIdentifier ().GetPriority ();
}

// ============================================================================

void STP_SetPortPriority (BRIDGE* bridge, int portIndex, int treeIndex, unsigned char portPriority, unsigned long timestampMilliseconds)
{
	// See table 13-3 on page 348 of 802.1Q-2011
	// See 13.25.32 in 802.1Q-2011
	
	assert ((portPriority % 16) == 0);

	assert ((portIndex >= 0) && (portIndex < bridge->portCount));
	assert ((treeIndex >= 0) && (treeIndex < bridge->treeCount));

	LOG (bridge, -1, -1, "{T}: Setting port priority: port {S} tree {TN} prio = {D}...\r\n",
		 timestampMilliseconds,
		 bridge->ports [portIndex]->debugName,
		 treeIndex,
		 portPriority);

	unsigned short* pid = &bridge->ports [portIndex]->trees [treeIndex]->portId;
	unsigned short newPortId = ((unsigned short) portPriority << 8) | (*pid & 0x0FFF);
	*pid = newPortId;

	// It would make sense that stuff is recomputed also when the port priority in the portId variable
	// is changed (as it is recomputed for the bridge priority), but either the spec does not mention this, or I'm not seeing it.
	// Anyway, information about the new port priority can only be propagated by such a recomputation, so let's do that.
	if (bridge->started)
		RecomputePrioritiesAndPortRoles (bridge, treeIndex);

	LOG (bridge, -1, -1, "------------------------------------\r\n");
	FLUSH_LOG (bridge);
}

unsigned char STP_GetPortPriority (BRIDGE* bridge, int portIndex, int treeIndex)
{
	// See 13.25.32 in 802.1Q-2011
	
	assert ((portIndex >= 0) && (portIndex < bridge->portCount));
	assert ((treeIndex >= 0) && (treeIndex < bridge->treeCount));

	unsigned short portId = bridge->ports [portIndex]->trees [treeIndex]->portId & 0xff00;
	return (portId >> 12) << 4;
}

// ============================================================================
